'''
Created on Feb 14, 2010

@author: oabalbin
'''

import numpy as np
import scipy as sp
from scipy import stats
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import matplotlib.pylab as pla
import matplotlib.cm as cm 

import rpy2.robjects as robjects
import rpy2.robjects.numpy2ri
from rpy2.robjects.packages import importr
from pylab import *


import signatures.parsers.read_gene_lists as gpm
import signatures.preprocess.tools as tl


def plot_scatter(mat1,mat2, corr, pli, folderpath,geneName=[]):
    """
    outa[0,0],outa[0,1],outa[0,3],outa[0,4], outa[0,5] = i (mat1 row index), k(factor number), pval, pcor, bpval, int(cmpname[i]) 
    (drug index in the NCI database).
    """    
    '''Linear regression'''
    
    #mat2 = 2*(mat2/float((np.max(mat2) - np.min(mat2))))
    xmin,xmax,ymin,ymax = np.min(mat1),np.max(mat1),np.min(mat2),np.max(mat2)
    
    # Linear regression    
    x = np.linspace(xmin,xmax,mat1.size)
    (ar,br) = np.polyfit(x,mat2,1)
    yn=np.polyval([ar,br],x)
    
    mat1sd = np.std(mat1)
    mat2sd = np.std(mat2)

    plt.xlim( (xmin-mat1sd, xmax+mat1sd) )
    plt.ylim( (ymin-0.5, ymax+0.5) )
    #plt.title('Linear regression with polifit()')
    #plt.plot(x,yn,'b-') #label='linear regression'
    plt.plot(mat1,mat2, 'ko')
    #legend(loc='best')
        
    plt.ylabel('Pathway activity (score)')
    plt.xlabel('Drug sensitivity (GI50) pval=%.4f' % round(corr[4],4) + ' R=%.2f' % round(corr[3],2))
    #xlim( (xmin, xmax) )     
    #plt.show()
    # Remember that not intercept is in the factors so the actual number of the factors with BFRM
    # Corresponds to j+1. corr[5] corresponds to the index of the drug in the NCI drug database.
    
    
    if geneName:
        pltname=folderpath+'plt_'+str(corr[5])+'_'+str(geneName[int(corr[1])])+'_'+str(pli)+'_.png'
    else:
        pltname=folderpath+'plt_'+str(corr[5])+'_'+str(corr[1] + 1)+'.png'

    plt.savefig(pltname)
    # Change to use with the general progam
    #plt.savefig(folderpath+'plt_'+str(corr[5])+'_'+str(corr[1])+'_'+'.png')
    plt.close()
    


def plot_scatter2(mat1,mat2, corr, folderpath):
    
    plt.scatter(mat1, mat2, c='k', marker='o')
    
    plt.savefig(folderpath+'plt_'+str(corr[0])+'_'+str(corr[1])+'_'+'.png')
    # Change to use with the general progam
    #plt.savefig(folderpath+'plt_'+str(corr[5])+'_'+str(corr[1])+'_'+'.png')
    plt.close()



def plot_drugvsact_2(mat1,mat2, corrmat, folderpath,cmpname):
    """
    """
    print corrmat.shape
    
    outfile = open(folderpath+'compound_indexes.txt','w') 
    
    cmpname = np.array(map(int,cmpname))
    index = np.array(range(mat1.shape[0]))
    
    for i in range(corrmat.shape[0]):
        r = index[cmpname==int(corrmat[i,0])][0]
        print str(int(corrmat[i,0]))+'\t'+str(corrmat[i,1])+'\t'+str(r)
        outfile.write(str(int(corrmat[i,0]))+'\t'+str(corrmat[i,1])+'\t'+str(r)+'\t'+str(corrmat[i,2])+'\t'+str(corrmat[i,3]) +'\n')

        #if i < 20:
        #plot_scatter(mat1[corrmat[i,0],:],mat2[corrmat[i,1],:], corrmat[i,:], folderpath)
        plot_scatter(mat1[r,:],mat2[corrmat[i,1],:], corrmat[i,:], folderpath)
        #else:
        #    break
    outfile.close()
    

def write_drugpath_mat(mat1, mat2, corrmat0, folderpath, samplenames, masked=False, cmpname=[], geneNames=[]):
    """
    Write a matrix of drug-module correlations. It could be used for downstream clustering analysis.
    corrmat (outa): matrix of correlation between GI50 and pathactivity.
    outa[0,0],outa[0,1],outa[0,3],outa[0,4], outa[0,5] = i (mat1 row index), k(factor number), pval, pcor, bpval, int(cmpname[i]) 
    (drug index in the NCI database).
    outa.shape[0] is smaller than mat1.shape[0]. It is a subset ot it.
    """
    
    #outdrugmat = open(folderpath+'drug_mat_sensit.txt','w')
    # drug indices
    #cmpname = np.array(map(int,cmpname))
    
    # number of rows in the drug-module correlation matrix.
    cmpindex = np.array(range(corrmat0.shape[0]))
    
    ## sort the correlation matrix according to the correlation value.
    sort_ind = list(np.argsort(corrmat0[:,3]))
    sort_ind.reverse()
    sort_ind = np.array(sort_ind)
    corrmat = corrmat0[sort_ind,:]
    
    
    # Factors for which there was correlation with some drug
    factors = np.array(map(int,set(corrmat[:,1])))
    print factors
    
    for j in factors:
        # find rows corresponding to factor j in the correlation matrix
        # Remember that not intercept is in the factors so the actual number of the factors with BFRM
        # Corresponds to j+1
        tf = j+1
        if geneNames:
            tf = geneNames[j]
        print tf
        # the index for the row of each correlated compound in the original mat1. 
        # This number was stored in the tl.calc_correlation_mats
        findx = corrmat[corrmat[:,1]==j,0]
        cindx = cmpindex[corrmat[:,1]==j]
        
        filename = folderpath+'fact_'+str(tf)+'_drug_mat_sensit.txt'
        outdrugmat = open(filename,'w')
        
        # sort the module activity in increasing order
        if masked:
            act_ind = np.ma.argsort(mat2[j,:])
        else:
            act_ind = np.argsort(mat2[j,:])
        
        linef = map(str,list(mat2[j,act_ind]))
        
        # write
        samplenames = np.array(samplenames)
        if cmpname:
            outdrugmat.write(str('compound')+'\t'+str('cmpname')+'\t'+str('factor')+'\t'+str('corr')+'\t'+str('pval')+'\t'+",".join(list(samplenames[act_ind])).replace(',','\t')+'\n')
            outdrugmat.write(str('m'+str(tf))+'\t'+str(tf)+'\t'+str('NA')+'\t'+str('NA')+'\t'+str('NA')+'\t'+",".join(linef).replace(',','\t')+'\n')
        else:
            outdrugmat.write(str('compound')+'\t'+str('factor')+'\t'+str('corr')+'\t'+str('pval')+'\t'+",".join(list(samplenames[act_ind])).replace(',','\t')+'\n')
            outdrugmat.write(str('m'+str(tf))+'\t'+str(tf)+'\t'+str('NA')+'\t'+str('NA')+'\t'+",".join(linef).replace(',','\t')+'\n')
        
        for k,i in zip(findx,cindx):
            # print ordered GI50 values according to the path_activity. for each correlated compound.
            print k, corrmat[i,0], corrmat[i,5],corrmat[i,1]
            line = map(str,list(mat1[k,act_ind]))
            if cmpname:
                outdrugmat.write(str(corrmat[i,5])+'\t'+str(cmpname[int(corrmat[i,0])])+'\t'+str(tf)+'\t'+str(corrmat[i,3])+'\t'+str(corrmat[i,4])+'\t'+",".join(line).replace(',','\t')+'\n')
            else:
                outdrugmat.write(str(corrmat[i,5])+'\t'+str(tf)+'\t'+str(corrmat[i,3])+'\t'+str(corrmat[i,4])+'\t'+",".join(line).replace(',','\t')+'\n')
    
    

def plot_drugvsact(mat1,mat2, corrmat, folderpath,pairwisemasking=False, geneNames=[]):
    """
    mat1: matrix of the drug sensitivities.
    mat2: matrix of to pathway activities.
    corrmat (outa): matrix of correlation between GI50 and pathactivity.
    outa[0,0],outa[0,1],outa[0,3],outa[0,4], outa[0,5] = i (mat1 row index), k(factor number), pval, pcor, bpval, int(cmpname[i]) 
    (drug index in the NCI database).
    """
    
    outfile = open(folderpath+'compound_indexes.txt','w') 
    
    for i in range(corrmat.shape[0]): #(range(10)):#range(corrmat.shape[0]):
        outfile.write(str(int(corrmat[i,0]))+'\t'+str(corrmat[i,1])+'\t'+str(corrmat[i,2])+'\t'+str(corrmat[i,3]) +'\t'+str(corrmat[i,4])+'\t'+str(corrmat[i,5]) +'\n')
        GI50 = mat1[corrmat[i,0],:]
        path = mat2[corrmat[i,1],:]
        
        if pairwisemasking:
            GI50, path = tl.calc_parwise_masking(GI50,path)
            
        if abs(corrmat[i,3]) > 0.35:
            plot_scatter(GI50, path, corrmat[i,:], i, folderpath, geneNames)



def plot_heatmap(pltname,expMat, row_names, col_names):
    """
    It plots a heat map for a given expression value matrix.
    """
    
    """
    def patient_colour(patient_id) :
    assert patient_id in mol_biol, \
           "Patient ID of '%s' not in list!" % patient_id
    if mol_biol[patient_id] == "ALL1/AF4" :
        return "#FF0000" # Red
    else :
        return "#0000FF" # Blue
    patient_colours = map(patient_colour, col_names)    
    """
    row_names=list(row_names)
    col_names =list(col_names)
    grdevices = importr('grDevices')
    grdevices.png(file=pltname, width=600, height=589)
    
    r = robjects.r
    r.png(pltname, width=600, height=589)
    r.heatmap(expMat,cexRow=0.5,labRow=row_names, labCol=col_names)
    
    grdevices.dev_off()


def plot_heatmap_mod(pltname,expMat, row_names, col_names, samples_status):
    """
    It plots a heat map for a given expression value matrix.
    """
    
    """
    def patient_colour(patient_id) :
    assert patient_id in mol_biol, \
           "Patient ID of '%s' not in list!" % patient_id
    if mol_biol[patient_id] == "ALL1/AF4" :
        return "#FF0000" # Red
    else :
        return "#0000FF" # Blue
    patient_colours = map(patient_colour, col_names)    
    """
    color_for_samples = np.array(map(sample_color,samples_status))
    print color_for_samples
    print color_for_samples.shape, len(col_names)
    
    
    #set_default_mode(BASIC_CONVERSION)
    
    row_names=list(row_names)
    col_names =list(col_names)
    xlabel_plot = 'Tumor samples'
    ylabel_plot = 'Modules'
    
    grdevices = importr('grDevices')
    grdevices.png(file=pltname, width=600, height=589)
    
    r = robjects.r
    r.png(pltname, width=600, height=589)
    r.library("grDevices")
    r.library("gplots")
    #r.par("cex.axis=1.2")
    #topocolors = r['topo.colors']
    #r.heatmap(expMat,cexRow=0.5,labRow=row_names, labCol=col_names, ColSideColors = color_for_samples, col = topocolors(50))

    heatmap2 = r['heatmap.2']
    
    heatmap2(expMat,
            labRow=row_names,
            scale="none",
            labCol=col_names,
            dendrogram = "column",
            ColSideColors=color_for_samples,
            col=r.bluered(75),
            key=True,
            keysize=1,
            symkey=False,
            density_info="none",
            trace="none",
            cexRow=0.05,
            xlab=xlabel_plot, 
            ylab=ylabel_plot,
            margins=r.c(3,3)
            )
        
    grdevices.dev_off()
    
    
    """    
    heatmap(x, Rowv=NULL, Colv=if(symm)"Rowv" else NULL,
        distfun = dist, hclustfun = hclust,
        reorderfun = function(d,w) reorder(d,w),
        add.expr, symm = FALSE, revC = identical(Colv, "Rowv"),
        scale=c("row", "column", "none"), na.rm = TRUE,
        margins = c(5, 5), ColSideColors, RowSideColors,
        cexRow = 0.2 + 1/log10(nr), cexCol = 0.2 + 1/log10(nc),
        labRow = NULL, labCol = NULL, main = NULL,
        xlab = NULL, ylab = NULL,
        keep.dendro = FALSE, verbose = getOption("verbose"), ...)
        
        
    heatmap.2 (x,

           # dendrogram control
           Rowv = TRUE,
           Colv=if(symm)"Rowv" else TRUE,
           distfun = dist,
           hclustfun = hclust,
           dendrogram = c("both","row","column","none"),
           symm = FALSE,

           # data scaling
           scale = c("none","row", "column"),
           na.rm=TRUE,

           # image plot
           revC = identical(Colv, "Rowv"),
           add.expr,

           # mapping data to colors
           breaks,
           symbreaks=min(x < 0, na.rm=TRUE) || scale!="none",

           # colors
           col="heat.colors",

           # block sepration
           colsep,
           rowsep,
           sepcolor="white",
           sepwidth=c(0.05,0.05),

           # cell labeling
           cellnote,
           notecex=1.0,
           notecol="cyan",
           na.color=par("bg"),

           # level trace
           trace=c("column","row","both","none"),
           tracecol="cyan",
           hline=median(breaks),
           vline=median(breaks),
           linecol=tracecol,

           # Row/Column Labeling
           margins = c(5, 5),
           ColSideColors,
           RowSideColors,
           cexRow = 0.2 + 1/log10(nr),
           cexCol = 0.2 + 1/log10(nc),
           labRow = NULL,
           labCol = NULL,

           # color key + density info
           key = TRUE,
           keysize = 1.5,
           density.info=c("histogram","density","none"),
           denscol=tracecol,
           symkey = min(x < 0, na.rm=TRUE) || symbreaks,
           densadj = 0.25,

           # plot labels
           main = NULL,
           xlab = NULL,
           ylab = NULL,

           # plot layout
           lmat = NULL,
           lhei = NULL,
           lwid = NULL,

           # extras
           ...
           )


    
    """


    
def sample_color(ets_status):
    #color = {'ETS+':"#FF0000",'ETS-':"#0000FF"}[ets_status]
    color = {'ETS+':"yellow",'ETS-':"green"}[ets_status]
    return color 
    

def plot_heatmap_mod2(pltname,expMat, row_names, col_names, samples_status=[]):
    """
    """
    if samples_status:
        color_for_samples = np.array(map(sample_color,samples_status))
        print color_for_samples
        print color_for_samples.shape, len(col_names)
    
    
    #set_default_mode(BASIC_CONVERSION)
    
    row_names=list(row_names)
    col_names =list(col_names)
    nr = len(row_names)
    nc = len(col_names)
    #xlabel_plot = 'Tumor samples'
    #ylabel_plot = 'Modules'
        
    r = robjects.r
    r.library("grDevices")
    r.library("gplots")
    
    grdevices = importr('grDevices')
    grdevices.png(file=pltname, width=1000, height=1000)

    
    #r.par("cex.axis=1.2")
    #topocolors = r['topo.colors']
    #r.heatmap(expMat,cexRow=0.5,labRow=row_names, labCol=col_names, ColSideColors = color_for_samples, col = topocolors(50))

    heatmap2 = r['heatmap.2']
    
    heatmap2(expMat,
            labRow=row_names,
            scale="none",
            labCol=col_names,
            dendrogram = "none",
            col=r.bluered(75),
            key=False,
            #keysize=1,
            #symkey=False,
            density_info="none",
            trace="none",
            #cexRow=0.05,
            #xlab=xlabel_plot, 
            #ylab=ylabel_plot,
            #margins=r.c(3,3),
            margins = r.c(5, 5),
            #ColSideColors,
            #RowSideColors,
            cexRow = 0.5 + 1/np.log10(nr),
            cexCol = 0.5 + 1/np.log10(nc),
            #labRow = NULL,
            #labCol = NULL,

            
            
            )
        
    grdevices.dev_off()





        
def density_plot(pltname,expMat):
    """
    Plots the density function
    """
    
    grdevices = importr('grDevices')
    grdevices.png(file=pltname, width=600, height=589)
    
    r = robjects.r
    r.png(pltname, width=600, height=589)
    
    for i in range(expMat.shape[1]):
        
        sp_i = expMat[:,i]
        #sp_i_density = r.density(sp_i,kernel="gaussian")
        if i==0:
            r.plot(r.density(sp_i,kernel="gaussian"))
        else:
            r.lines(r.density(sp_i,kernel="gaussian"))
        
    grdevices.dev_off()
    

def plot_correlation_matrix_py(A):
    """
    """
    #A = np.random.randint(10, 100, 100).reshape(10, 10)
    #mask =  np.tri(A.shape[0],A.shape[1], k=-1)
    #A = np.ma.array(A, mask=mask) # mask out the lower triangle
    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    cmap = cm.get_cmap('jet', 10) # jet doesn't have white color
    cmap.set_bad('w') # default value is 'k'
    ax1.imshow(A, interpolation="nearest", cmap=cmap)
    ax1.grid(True)
    plt.show()

    


#numpy.loadtxt(fname, dtype=<type 'float'>, comments='#', delimiter=None, converters=None, skiprows=0, usecols=None, unpack=False)


'''
#numpy.loadtxt(fname, dtype=<type 'float'>, comments='#', delimiter=None, converters=None, skiprows=0, usecols=None, unpack=False)
mat2 = np.loadtxt('/home/oabalbin/projects/networks/output/2010_03_20_21_13/bfrm/mF.txt',dtype=float, delimiter='\t',skiprows=(0))
ga = gpm.geneparser()
samples = ga.list_of_names_in_line(open('/home/oabalbin/projects/networks/output/2010_03_20_21_13/bfrm/sample_test'))
dictsamples = []
for i,sp in enumerate(samples):
    if i < 40:
        dictsamples.append('ETS+')
    else:
        dictsamples.append('ETS-')
     
     
modules= np.array(range(mat2.shape[0]))
pltname = '/home/oabalbin/projects/networks/output/2010_03_20_21_13/sdv/heat_map_test.png'

plot_heatmap_mod(pltname, mat2, modules, samples, dictsamples)
'''

    
    
    